#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : 应用案例.py
# @Author: HCY
# @Date  : 2025/5/17
# @Desc  : 讯飞星火知识库服务应用案例

import os
import time
import json
import gradio as gr
from typing import List, Dict
import tempfile
import shutil

# SDK引入模型
from dwspark.config import Config
from dwspark.models import KnowledgeBase
# 日志
from loguru import logger

# 加载系统环境变量：SPARKAI_UID、SPARKAI_APP_ID、SPARKAI_API_KEY、SPARKAI_API_SECRET
config = Config(appid='4095e411', apisecret='ZGRiMWNkM2EzYzdmZWNlZTI2YjMyYzdh', apikey='5bed18a456d46392e0ffa240fee8990c')

# 创建星火知识库实例
kb = KnowledgeBase(config)

# 存储当前会话的知识库ID和文件ID
current_repo_id = None
uploaded_files = []

def create_repository(repo_name: str, repo_desc: str, repo_tags: str):
    """
    创建知识库
    :param repo_name: 知识库名称
    :param repo_desc: 知识库描述
    :param repo_tags: 知识库标签，用逗号分隔
    :return: 知识库ID和状态信息
    """
    global current_repo_id
    try:
        # 获取现有知识库列表
        url = f"{kb.REPO_API_URL}/repo/list"
        headers = kb._get_headers()
        body = {
            "currentPage": 1,
            "pageSize": 100  # 设置较大的页面大小以获取更多知识库
        }

        import requests
        response = requests.post(url, json=body, headers=headers)
        result = kb._handle_response(response, "获取知识库列表")
        
        if result and "data" in result:
            repos = result["data"].get("rows", [])
            # 检查是否存在同名知识库
            for repo in repos:
                if repo.get("repoName") == repo_name:
                    return f"创建知识库失败：已存在同名知识库 '{repo_name}'", "未创建知识库"
        
        # 创建新知识库
        current_repo_id = kb.create_repository(repo_name, repo_desc, repo_tags)
        return f"创建知识库成功！知识库ID: {current_repo_id}", f"已创建知识库: {repo_name}"
    except Exception as e:
        logger.error(f"创建知识库失败: {e}")
        return f"创建知识库失败: {str(e)}", "未创建知识库"

def upload_document(file_obj):
    """
    上传文档到知识库
    :param file_obj: 上传的文件对象
    :return: 文件ID和状态信息
    """
    global current_repo_id, uploaded_files
    
    if file_obj is None:
        return "请选择要上传的文件", f"已上传文件: {len(uploaded_files)}个"
    
    try:
        # 获取文件路径
        file_path = file_obj.name
        
        # 检查文件是否存在
        if not os.path.exists(file_path):
            return f"文件不存在: {file_path}", f"已上传文件: {len(uploaded_files)}个"
        
        # 上传文档
        file_id = kb.upload_document(file_path)
        
        # 如果已创建知识库，则将文档添加到知识库
        if current_repo_id:
            url = f"{kb.REPO_API_URL}/repo/file/add"
            headers = kb._get_headers()
            body = {
                "repoId": current_repo_id,
                "fileIds": [file_id]
            }
            
            import requests
            response = requests.post(url, json=body, headers=headers)
            result = kb._handle_response(response, "添加文档到知识库")
            
            if result:
                uploaded_files.append({"id": file_id, "name": os.path.basename(file_path)})
                return f"文档上传并添加到知识库成功！文件ID: {file_id}", f"已上传文件: {len(uploaded_files)}个"
        else:
            uploaded_files.append({"id": file_id, "name": os.path.basename(file_path)})
            return f"文档上传成功！文件ID: {file_id}（未添加到知识库，请先创建知识库）", f"已上传文件: {len(uploaded_files)}个"
    except Exception as e:
        logger.error(f"上传文档失败: {e}")
        return f"上传文档失败: {str(e)}", f"已上传文件: {len(uploaded_files)}个"

def list_repositories():
    """
    获取知识库列表
    :return: 知识库列表信息
    """
    try:
        url = f"{kb.REPO_API_URL}/repo/list"
        headers = kb._get_headers()
        body = {
            "currentPage": 1,
            "pageSize": 10
        }

        import requests
        response = requests.post(url, json=body, headers=headers)
        result = kb._handle_response(response, "获取知识库列表")
        
        if result and "data" in result:
            repos = result["data"].get("rows", [])
            repo_info = ""
            for i, repo in enumerate(repos):
                repo_info += f"{i+1}. {repo.get('repoName')} (ID: {repo.get('repoId')})\n"
            
            if not repo_info:
                return "暂无知识库"
            return repo_info
        else:
            return "获取知识库列表失败"
    except Exception as e:
        logger.error(f"获取知识库列表失败: {e}")
        return f"获取知识库列表失败: {str(e)}"

def list_uploaded_files():
    """
    获取已上传文件列表
    :return: 文件列表信息
    """
    global uploaded_files
    try:
        file_info = ""
        for i, file in enumerate(uploaded_files):
            file_info += f"{i+1}. {file.get('name')} (ID: {file.get('id')})\n"
        
        if not file_info:
            return "暂无上传文件"
        return file_info
    except Exception as e:
        logger.error(f"获取文件列表失败: {e}")
        return f"获取文件列表失败: {str(e)}"

def select_repository(repo_id: str):
    """
    选择知识库
    :param repo_id: 知识库ID
    :return: 状态信息
    """
    global current_repo_id
    try:
        current_repo_id = repo_id
        return f"已选择知识库，ID: {repo_id}", f"当前知识库ID: {repo_id}"
    except Exception as e:
        logger.error(f"选择知识库失败: {e}")
        return f"选择知识库失败: {str(e)}", "未选择知识库"

def create_test_file():
    """
    创建测试文件
    :return: 文件路径
    """
    try:
        # 创建临时目录
        temp_dir = tempfile.mkdtemp()
        file_path = os.path.join(temp_dir, "test_document.txt")
        
        # 写入测试内容
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write("这是一个测试文档，用于演示星火知识库API的使用。\n")
            f.write("星火大模型是讯飞推出的大语言模型，具有强大的自然语言处理能力。\n")
            f.write("知识库可以帮助用户存储和检索各种文档资料，提高信息获取效率。")
        
        return f"已创建测试文件: {file_path}", file_path
    except Exception as e:
        logger.error(f"创建测试文件失败: {e}")
        return f"创建测试文件失败: {str(e)}", None

def upload_test_file(file_path):
    """
    上传测试文件
    :param file_path: 测试文件路径
    :return: 上传结果
    """
    if not file_path:
        return "无测试文件路径", f"已上传文件: {len(uploaded_files)}个"
    
    try:
        # 上传文档
        file_id = kb.upload_document(file_path)
        
        # 等待文档处理完成
        max_retries = 5
        retry_interval = 2  # 秒
        
        for _ in range(max_retries):
            # 获取文档详情，检查处理状态
            doc_info = kb.get_document_info(file_id)
            if doc_info and doc_info.get('status') == 'FINISHED':
                break
            time.sleep(retry_interval)
        
        # 如果已创建知识库，则将文档添加到知识库
        if current_repo_id:
            url = f"{kb.REPO_API_URL}/repo/file/add"
            headers = kb._get_headers()
            body = {
                "repoId": current_repo_id,
                "fileIds": [file_id]
            }
            
            import requests
            response = requests.post(url, json=body, headers=headers)
            result = kb._handle_response(response, "添加文档到知识库")
            
            if result:
                # 等待文档在知识库中处理完成
                time.sleep(5)  # 给予足够的处理时间
                uploaded_files.append({"id": file_id, "name": os.path.basename(file_path)})
                return f"测试文档上传并添加到知识库成功！文件ID: {file_id}", f"已上传文件: {len(uploaded_files)}个"
        else:
            uploaded_files.append({"id": file_id, "name": os.path.basename(file_path)})
            return f"测试文档上传成功！文件ID: {file_id}（未添加到知识库，请先创建知识库）", f"已上传文件: {len(uploaded_files)}个"
    except Exception as e:
        logger.error(f"上传测试文档失败: {e}")
        return f"上传测试文档失败: {str(e)}", f"已上传文件: {len(uploaded_files)}个"

def chat_with_kb(query: str, chat_history: List):
    """
    与知识库对话
    :param query: 用户问题
    :param chat_history: 历史对话
    :return: 更新后的对话历史
    """
    global current_repo_id
    if not current_repo_id:
        chat_history.append((query, "请先创建或选择知识库！"))
        return "", chat_history
    
    try:
        # 添加用户问题到对话历史
        chat_history.append((query, ""))
        
        # 定义回调函数
        # 用于累积消息文本
        response_text = []
        response_complete = False
        
        def on_message(data):
            nonlocal chat_history, response_text, response_complete
            logger.debug(f"收到消息: {data}")
            
            # 尝试解析字符串消息
            if isinstance(data, str):
                try:
                    data = json.loads(data)
                except json.JSONDecodeError:
                    logger.error(f"无法解析JSON消息: {data}")
                    # 如果无法解析为JSON，直接将其作为文本添加
                    response_text.append(data)
                    chat_history[-1] = (query, ''.join(response_text))
                    return

            # 检查不同的消息格式
            if isinstance(data, dict):
                # 检查是否有错误码
                if "code" in data and data["code"] != 0:
                    error_msg = data.get("message", "未知错误")
                    chat_history[-1] = (query, f"对话失败: {error_msg}")
                    response_complete = True
                    return
                
                # 处理content字段 - 知识库API返回的主要格式
                if 'content' in data:
                    content = data['content']
                    if content:
                        response_text.append(content)
                        chat_history[-1] = (query, ''.join(response_text))
                
                # 处理payload格式 - 备用格式
                elif 'payload' in data:
                    payload = data['payload']
                    if isinstance(payload, dict) and 'content' in payload:
                        content = payload['content']
                        if content:
                            response_text.append(content)
                            chat_history[-1] = (query, ''.join(response_text))
                
                # 检查是否完成
                if "status" in data and data["status"] == 2:
                    response_complete = True
                
        def on_error(error_msg):
            nonlocal chat_history, response_complete
            logger.error(f"对话失败: {error_msg}")
            chat_history[-1] = (query, f"对话失败: {error_msg}")
            response_complete = True
            
        def on_close():
            nonlocal response_complete
            logger.info("对话连接已关闭")
            response_complete = True
        
        # 调用知识库进行对话
        kb.chat(
            query=query,
            repo_id=current_repo_id,
            on_message=on_message,
            on_error=on_error,
            on_close=on_close
        )
        
        # 等待响应完成或超时
        max_wait_time = 30  # 最大等待时间，单位秒
        wait_start = time.time()
        
        while not response_complete and time.time() - wait_start < max_wait_time:
            time.sleep(0.5)
        
        # 如果响应为空，可能是回调没有被调用
        if not response_text:
            chat_history[-1] = (query, "未收到回复，请检查知识库是否包含相关信息或重试")
        
        # 返回清空输入框和更新后的对话历史
        return "", chat_history
    except Exception as e:
        logger.error(f"对话失败: {e}")
        chat_history[-1] = (query, f"对话失败: {str(e)}")
        return "", chat_history

# 构建Gradio界面
with gr.Blocks(title="星火知识库") as demo:
    gr.Markdown("# 讯飞星火知识库")
    
    # 创建状态显示组件
    status_display = gr.Textbox(label="知识库状态", value="未创建知识库", interactive=False)
    
    # 存储测试文件路径
    test_file_path = gr.State(None)
    
    with gr.Tab("知识库管理"):
        with gr.Row():
            with gr.Column():
                gr.Markdown("## 创建知识库")
                repo_name = gr.Textbox(label="知识库名称", placeholder="请输入知识库名称", value="测试知识库")
                repo_desc = gr.Textbox(label="知识库描述", placeholder="请输入知识库描述", value="这是一个测试知识库")
                repo_tags = gr.Textbox(label="知识库标签", placeholder="请输入标签，用逗号分隔", value="测试,文档,API")
                create_repo_btn = gr.Button("创建知识库")
                repo_result = gr.Textbox(label="创建结果", interactive=False)
                
                create_repo_btn.click(
                    fn=create_repository,
                    inputs=[repo_name, repo_desc, repo_tags],
                    outputs=[repo_result, status_display]
                )
            
            with gr.Column():
                gr.Markdown("## 知识库列表")
                list_repo_btn = gr.Button("刷新知识库列表")
                repo_list = gr.Textbox(label="知识库列表", interactive=False, lines=10)
                
                list_repo_btn.click(
                    fn=list_repositories,
                    inputs=[],
                    outputs=[repo_list]
                )
                
                gr.Markdown("## 选择知识库")
                repo_id_input = gr.Textbox(label="知识库ID", placeholder="请输入要选择的知识库ID")
                select_repo_btn = gr.Button("选择知识库")
                select_repo_result = gr.Textbox(label="选择结果", interactive=False)
                
                select_repo_btn.click(
                    fn=select_repository,
                    inputs=[repo_id_input],
                    outputs=[select_repo_result, status_display]
                )
    
    with gr.Tab("文档管理"):
        with gr.Row():
            with gr.Column():
                gr.Markdown("## 创建测试文档")
                create_test_btn = gr.Button("创建测试文档")
                test_result = gr.Textbox(label="创建结果", interactive=False)
                
                # 绑定创建测试文件按钮
                create_test_btn.click(
                    fn=create_test_file,
                    inputs=[],
                    outputs=[test_result, test_file_path]
                )
                
                gr.Markdown("## 上传测试文档")
                upload_test_btn = gr.Button("上传测试文档")
                upload_test_result = gr.Textbox(label="上传结果", interactive=False)
                file_count = gr.Textbox(label="文件计数", value="已上传文件: 0个", interactive=False)
                
                # 绑定上传测试文件按钮
                upload_test_btn.click(
                    fn=upload_test_file,
                    inputs=[test_file_path],
                    outputs=[upload_test_result, file_count]
                )
            
            with gr.Column():
                gr.Markdown("## 文件列表")
                list_files_btn = gr.Button("刷新文件列表")
                file_list = gr.Textbox(label="文件列表", interactive=False, lines=10)
                
                list_files_btn.click(
                    fn=list_uploaded_files,
                    inputs=[],
                    outputs=[file_list]
                )
    
    with gr.Tab("知识库对话"):
        gr.Markdown("## 与知识库对话")
        with gr.Row():
            with gr.Column(scale=2):
                chatbot = gr.Chatbot([], elem_id="chat-box", label="聊天历史")
                chat_input = gr.Textbox(label="输入问题", placeholder="请输入您想问的问题")
                chat_btn = gr.Button("发送")
        
        # 问题样例
        gr.Examples(["星火大模型有什么特点？", "知识库可以存储哪些类型的文档？", "如何使用星火知识库API？"], chat_input)
        
        # 绑定发送按钮点击事件
        chat_btn.click(
            fn=chat_with_kb,
            inputs=[chat_input, chatbot],
            outputs=[chat_input, chatbot]
        )
        
        # 绑定回车键发送
        chat_input.submit(
            fn=chat_with_kb,
            inputs=[chat_input, chatbot],
            outputs=[chat_input, chatbot]
        )

if __name__ == "__main__":
    demo.queue().launch()